package com.weijian.controller.admin;

import com.weijian.advisor.ChatMessageAdvisor;
import com.weijian.common.pojo.CommonResult;
import com.weijian.dal.vo.ConversationVO;
import com.weijian.service.ChatService;
import com.weijian.utils.ParseWordFileUtil;
import io.swagger.v3.oas.annotations.Operation;
import jakarta.annotation.security.PermitAll;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.util.List;

import static com.weijian.constants.SystemPrompt.DEEPPROMPT;
import static com.weijian.constants.SystemPrompt.USUALPROMPT;


@Slf4j
@RestController
@RequiredArgsConstructor
@RequestMapping("/conversation")
public class ChatController {


    final ChatMemory chatMemory;

    final ChatClient chatClient;

    final ChatService chatService;

    final VectorStore vectorStore;


    @PostMapping("/chat/create")
    @Operation(summary = "创建对话")
    @PermitAll
    public CommonResult<ConversationVO> create() {
        return CommonResult.success(chatService.create());
    }

    @PutMapping("/chat/edit")
    @Operation(summary = "修改对话")
    @PermitAll
    public CommonResult<Boolean> edit(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "title") String title) {
        chatService.edit(conversationId, title);
        return CommonResult.success(true);
    }

    @PostMapping("/chat")
    @Operation(summary = "通用对话")
    @PermitAll
    public Flux<String> chat(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "message") String message,
            @RequestParam(value = "files", required = false) List<MultipartFile> files) {

        if (files == null || files.isEmpty()) {
            return textChat(USUALPROMPT, message, conversationId);
        } else {
            return multiModalChatAndStore(USUALPROMPT, message, conversationId, files);
        }
    }

    @PostMapping("/deepchat")
    @Operation(summary = "深度对话")
    @PermitAll
    public Flux<String> deeepchat(
            @RequestParam(value = "conversationId") String conversationId,
            @RequestParam(value = "message") String message,
            @RequestParam(value = "files", required = false) List<MultipartFile> files) {
        if (files == null || files.isEmpty()) {
            return textChat(DEEPPROMPT, message, conversationId);
        } else {
            return multiModalChatAndStore(DEEPPROMPT, message, conversationId, files);
        }
    }


    @DeleteMapping("/chat/del")
    @Operation(summary = "删除对话")
    @PermitAll
    public CommonResult<Boolean> del(@RequestParam(value = "conversationId") String conversationId) {
        chatService.delete(conversationId);
        return CommonResult.success(true);
    }

    @GetMapping("/chat/list")
    @Operation(summary = "展示对话")
    @PermitAll
    public CommonResult<List<ConversationVO>> list() {
        return CommonResult.success(chatService.list());
    }

    @GetMapping("/chat/get")
    @Operation(summary = "获取对话")
    @PermitAll
    public CommonResult<ConversationVO> get(@RequestParam(value = "conversationId") String conversationId) {
        return CommonResult.success(chatService.get(conversationId));
    }

    @SneakyThrows
    private Flux<String> multiModalChatAndStore(String prompt, String message, String conversationId,
                                                List<MultipartFile> files) {
        return Flux.fromIterable(files)
                .flatMap(file -> {
                    try {
                        String fileName = file.getOriginalFilename();
                        String content = "";
                        if (fileName != null && fileName.endsWith(".docx")) {
                            content = ParseWordFileUtil.parseDocxFile(file);
                        } else if (fileName != null && fileName.endsWith(".doc")) {
                            content = ParseWordFileUtil.parseDocFile(file);
                        } else {
                            return Flux.error(new IllegalArgumentException("文件类型不支持，请上传word文件: " + fileName));
                        }
                        return Flux.just(content);
                    } catch (Exception e) {
                        return Flux.error(e);
                    }
                })
                .flatMap(content -> chatClient.prompt(prompt)
                        .user(p -> p.text(message).text(content))
                        .advisors(new ChatMessageAdvisor(chatMemory, conversationId))
                        .advisors(new QuestionAnswerAdvisor(vectorStore))
                        .stream()
                        .content())
                .doOnNext(System.out::print);
    }


    private Flux<String> textChat(String prompt, String message, String conversationId) {
        return chatClient.prompt(prompt)
                .user(message)
                .advisors(new ChatMessageAdvisor(chatMemory, conversationId))
                .advisors(new QuestionAnswerAdvisor(vectorStore))
                .stream()
                .content()
                .doOnNext(System.out::print);
    }

}
